package core.sql.udf;

import org.apache.spark.sql.*;
import org.apache.spark.sql.expressions.Aggregator;

/**
 *㚕：输入多个，返回一个，比如COUNT,SUM,AVG这种
 */
public class Spark03_UDAF {

    public static void main(String[] args) {
        SparkSession sparkSession = SparkSession
                .builder()
                .master("local[*]")
                .appName("Spark01_UDF")
                .getOrCreate();

        try {
            // 读取 JSON 数据
            Dataset<Row> json = sparkSession.read()
                    .option("mode", "PERMISSIVE") // 允许坏记录
                    .option("badRecordsPath", "data/bad_records") // 指定错误记录的输出路径
                    .json("data/user_sql.json");


            /**
             * SparkSql采用特殊的方法将UDAF转换成UDF使用
             */
            sparkSession.udf().register("avgAge", functions.udaf(
                    new avgAgeUdf(),Encoders.LONG()
            ));

            // 缓存解析后的结果
            json.cache();

            // 创建临时视图
            json.createOrReplaceTempView("user");

            // 示例：执行 SQL 查询
            Dataset<Row> result = sparkSession.sql("SELECT avgAge(age) as name FROM user");
            result.show(); // 显示结果

//            // 处理错误记录
//            Dataset<Row> corruptRecords = sparkSession.read().json("data/bad_records/*");
//            corruptRecords.show(); // 显示错误记录

        } catch (Exception e) {
            e.printStackTrace(); // 捕获并打印异常
        } finally {
            sparkSession.close();
        }
    }

}


/**
 * 自定义UDAF函数，实现年龄的平均值计算
 * 1、创建自定义的累
 * 2、继承org.apache.spark.sql.expressions.Aggregator
 * 3、设定参数，总共三个参数
 * 1.IN输入数据类型
 * 2.BUFF 缓冲区
 * 3.OUT输出数据类型
 * 4、重写方法
 */

class avgAgeUdf extends Aggregator<Long, AvgAfeBuff, Long> {


    /**
     * 初始化
     *
     * @return
     */
    @Override
    public AvgAfeBuff zero() {
        return new AvgAfeBuff(0L, 0L);
    }

    /**
     * 将数据年龄和BUFF进行聚合
     *
     * @param b
     * @param a
     * @return
     */
    @Override
    public AvgAfeBuff reduce(AvgAfeBuff b, Long a) {
        b.setCount(b.getCount() + 1);
        b.setSum(b.getSum() + a);
        return b;
    }

    /**
     * 合并缓冲区
     * @param b1
     * @param b2
     * @return
     */
    @Override
    public AvgAfeBuff merge(AvgAfeBuff b1, AvgAfeBuff b2) {
        b1.setCount(b1.getCount() + b2.getCount());
        b1.setSum(b1.getSum() + b2.getSum());
        return b1;
    }

    /**
     * 最终输出
     * @param reduction
     * @return
     */
    @Override
    public Long finish(AvgAfeBuff reduction) {
        return reduction.getSum()/reduction.getCount();
    }

    /**
     * BUFF的编码类型
     * @return
     */
    @Override
    public Encoder<AvgAfeBuff> bufferEncoder() {
        return Encoders.bean(AvgAfeBuff.class);
    }

    /**
     * 输出编码
     * @return
     */

    @Override
    public Encoder<Long> outputEncoder() {
        return Encoders.LONG();
    }
}


